package featureSelect;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;

import mlProject.DocModel;
import mlProject.DocModelReader;
import mlProject.DocModelWriter;

public class TrainFileSplit {
	public static void process(int numForTest) {
		ArrayList<DocModel> docs = DocModelReader.readFromFile("review.train");
		ArrayList<DocModel> newTrainDocs = new ArrayList<DocModel>();
		ArrayList<DocModel> newTestDocs = new ArrayList<DocModel>();
		Map<String, ArrayList<DocModel>> newTrainDocsSplitByDomain = new HashMap<String, ArrayList<DocModel>>();
		for (Integer i = 0; i < 4; i++) {
			ArrayList<DocModel> domainDocs=new ArrayList<DocModel>();
			for (int j = 1; j <= 1000; j++) {
				int index = i * 1000 + j-1;
				if (j >= 1 && j <= 500) {
					if (j <= 500 - numForTest) {
						newTrainDocs.add(docs.get(index));
						domainDocs.add(docs.get(index));
					} else {
						newTestDocs.add(docs.get(index));
					}
				} else {
					if (j >= 501 && j <= 1000 - numForTest) {
						newTrainDocs.add(docs.get(index));
						domainDocs.add(docs.get(index));
					} else {
						newTestDocs.add(docs.get(index));
					}
				}
			}
			newTrainDocsSplitByDomain.put(i.toString(), domainDocs);
		}
		//write to file
		String path="src/main/resources/";
		DocModelWriter.writeToFile(newTrainDocs, path+"newTrain.txt");
		DocModelWriter.writeToFile(newTestDocs,path+"newTest.txt");
		DocModelWriter.writeToFile(newTrainDocsSplitByDomain.get("0"),path+"book.txt");
		DocModelWriter.writeToFile(newTrainDocsSplitByDomain.get("1"),path+"DVD.txt");
		DocModelWriter.writeToFile(newTrainDocsSplitByDomain.get("2"),path+"electronic.txt");
		DocModelWriter.writeToFile(newTrainDocsSplitByDomain.get("3"),path+"kitchen.txt");
	}
	public static void main(String[] args){
		process(100);
	}
}
